import numpy as np
import matplotlib.pyplot as plt
import os
from matplotlib.colors import LogNorm

fname = 'fig4.npz'

# load npz file
data = np.load(fname, allow_pickle=True)
print(data.files)
n_up = data['n_up']
nu_up = data['nu_up']
B_up = data['B_up']
Rxy_up = data['Rxy_up']
Rxy_upasym = data['Rxy_upasym']
n_down = data['n_down']
nu_down = data['nu_down']
B_down = data['B_down']
Rxy_down = data['Rxy_down']
Rxy_downasym = data['Rxy_downasym']
Vbg = data['Vbg']
Vtg = data['Vtg'] 

n0 = data['n0'] # density at charge neutrality point
D0 = data['D0'] 
ns = data['ns'] # density at 4 electron per moire unit cell
device_name = data['device_name']
back_gate_thickness_nm = data['back_gate_thickness_nm']
top_gate_thickness_nm = data['top_gate_thickness_nm']

# plot raw data
fig, ax = plt.subplots(1, 2, figsize=(8, 3))
plt.subplots_adjust(wspace=0.4)
plt.subplots_adjust(bottom=0.2)
fig.suptitle('Fig. 4 raw data')

im1 = ax[0].pcolormesh(nu_up, B_up, Rxy_up/1e3, cmap='coolwarm', )
ax[0].set_xlabel('Density ($10^{12}$ cm$^{-2}$)')
ax[0].set_ylabel('$B$ (T)')
cbar = fig.colorbar(im1, ax=ax[0])
cbar.set_label(r'$R_{xy}\uparrow$ (k$\Omega$)')

im2 = ax[1].pcolormesh(nu_down, B_down, Rxy_down/1e3, cmap='coolwarm', )
ax[0].set_xlabel('Density ($10^{12}$ cm$^{-2}$)')
ax[0].set_ylabel('$B$ (T)')
cbar = fig.colorbar(im2, ax=ax[1])
cbar.set_label(r'$R_{xy}\downarrow$ (k$\Omega$)')
plt.show()

# plot symmetrized/antisymmetrized data
fig, ax = plt.subplots(1, 2, figsize=(8, 3))
plt.subplots_adjust(wspace=0.4)
plt.subplots_adjust(bottom=0.2)
fig.suptitle('Fig. 4 symmetrized/antisymmetrized data')

im1 = ax[0].pcolormesh(nu_up, B_up, Rxy_upasym/1e3, cmap='coolwarm', vmin=-1.2, vmax=1.2)
ax[0].set_xlabel('Density ($10^{12}$ cm$^{-2}$)')
ax[0].set_ylabel('$B$ (T)')
cbar = fig.colorbar(im1, ax=ax[0])
cbar.set_label(r'$R_{xy}\uparrow$ (k$\Omega$)')

im2 = ax[1].pcolormesh(nu_down, B_down, Rxy_downasym/1e3, cmap='coolwarm', vmin=-1.2, vmax=1.2)
ax[0].set_xlabel('Density ($10^{12}$ cm$^{-2}$)')
ax[0].set_ylabel('$B$ (T)')
cbar = fig.colorbar(im2, ax=ax[1])
cbar.set_label(r'$R_{xy}\downarrow$ (k$\Omega$)')
plt.show()
